import numpy as np
import itertools
import pickle
import scipy.spatial

class ToNumpy:
    def __call__(self, image, params=None):
        image = np.asarray(image)
        if len(image.shape) > 2:
            image = image[:,:,:3]
        return image

class FixMaxMinNorm:
    def __call__(self, image, params=None):
        if params is not None:
            pixel_max = params.get('max', None)
            pixel_min = params.get('min', None)
            if len(image.shape) == 3:
                pixel_max = np.array(pixel_max).reshape(1,1,-1)
                pixel_min = np.array(pixel_min).reshape(1,1,-1)
            else:
                pixel_max = np.array(pixel_max).reshape(1,1)
                pixel_min = np.array(pixel_min).reshape(1,1)
            image = (image - pixel_min) / (pixel_max - pixel_min)
        return image

class MaxMinNorm:
    def __call__(self, image, params=None):
        max_pixel = image.max(axis=0, keepdims=True).max(axis=1, keepdims=True)
        min_pixel = image.min(axis=0, keepdims=True).min(axis=1, keepdims=True)
        image = (image - min_pixel) / np.where(max_pixel == min_pixel, 1, max_pixel - min_pixel)
        return image
        
class ZScoreNorm:
    def __call__(self, image, params=None):
        if params is not None:
            pixel_mean = params.get('mean', None)
            pixel_std = params.get('std', None)
            if len(image.shape) == 3:
                pixel_mean = np.array(pixel_mean).reshape(1,1,-1)
                pixel_std = np.array(pixel_std).reshape(1,1,-1)
            else:
                pixel_mean = np.array(pixel_mean).reshape(1,1)
                pixel_std = np.array(pixel_std).reshape(1,1)
            image = (image - pixel_mean) / pixel_std
        return image        

class ExtractOneChannel:
    def __init__(self, channel_id):
        self.channel_id = channel_id

    def __call__(self, image, params=None):
        if len(image.shape) == 3:
            image = image[:, :, self.channel_id]
        return image

class PILImageCrop:
    def __call__(self, image, params=None):
        crop_box = params.get('crop', None) if params is not None else None
        if crop_box is not None:
            crop_box = [int(i) for i in crop_box]
            image = image.crop(crop_box)
        return image

class Repeat:
    def __call__(self, image, params=None):
        repeat_params = params.get('repeat', None) if params is not None else None
        if repeat_params is not None:
            u_repeat, v_repeat = repeat_params
            shape = list(image.shape)
            h, w = shape[:2]
            shape[0] *= v_repeat
            shape[1] *= u_repeat
            new_image = np.empty(shape)
            for u, v in itertools.product(list(range(u_repeat)), list(range(v_repeat))):
                new_image[v*h:(v+1)*h, u*w:(u+1)*w] = image
            image = new_image
        return image

class NearestNeighbor:
    def __init__(self, centers_path):
        self.centers_path = centers_path
        self.centers = None
        self.kdtree = None

    def __call__(self, image, params=None):
        if self.centers is None:
            with open(self.centers_path, 'rb') as f:
                self.centers = pickle.load(f)
            self.kdtree = scipy.spatial.cKDTree(self.centers)
        
        h, w, c = image.shape
        image = self.kdtree.query(image.reshape(-1, c), k=1)[1].reshape(h, w).astype(np.uint8)
        return image
